import os
import json
import argparse
import tensorrt as trt
import os
os.environ['TORCH_CUDA_ARCH_LIST'] = '7.5'
import sys
import time
import torch
import numpy as np
from typing import Iterable, Tuple

# === 路径加入工程根目录 ===
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from trt_ppq_deploy import build_trt_engine_with_ppq_json

ONNX_PATH        = os.path.join(PROJECT_ROOT, 'engine/quantization_onnx/yolov11n_int8_ALLConcat.onnx')
PPQ_JSON_PATH    = os.path.join(PROJECT_ROOT, 'engine/quantization_onnx/yolov11n_int8_ALLConcat.json')
ENGINE_OUTPUT_PATH = os.path.join(PROJECT_ROOT, 'engine/yolov11n_int8_ALLConcat.engine')

TRT_LOGGER = trt.Logger()

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

def GiB(val):
    return val * 1 << 30

def json_load(filename):
    with open(filename) as json_file:
        data = json.load(json_file)
    return data

def setDynamicRange(network, json_file):
    """Sets ranges for network layers."""
    quant_param_json = json_load(json_file)
    act_quant = quant_param_json["act_quant_info"]

    for i in range(network.num_inputs):
        input_tensor = network.get_input(i)
        if act_quant.__contains__(input_tensor.name):
            value = act_quant[input_tensor.name]
            tensor_max = abs(value)
            tensor_min = -abs(value)
            input_tensor.set_dynamic_range(tensor_min, tensor_max)


    for i in range(network.num_layers):
        layer = network.get_layer(i)

        for output_index in range(layer.num_outputs):
            tensor = layer.get_output(output_index)

            if act_quant.__contains__(tensor.name):
                value = act_quant[tensor.name]
                tensor_max = abs(value)
                tensor_min = -abs(value)
                tensor.set_dynamic_range(tensor_min, tensor_max)
            else:
                print("\033[1;32m%s\033[0m" % tensor.name)


def build_engine(onnx_file, json_file, engine_file):
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(EXPLICIT_BATCH)
    
    config = builder.create_builder_config()

    # If it is a dynamic onnx model , you need to add the following.
    # profile = builder.create_optimization_profile()
    # profile.set_shape("input_name", (batch, channels, min_h, min_w), (batch, channels, opt_h, opt_w), (batch, channels, max_h, max_w)) 
    # config.add_optimization_profile(profile)
    

    parser = trt.OnnxParser(network, TRT_LOGGER)
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, GiB(1))

    if not os.path.exists(onnx_file):
        quit('ONNX file {} not found'.format(onnx_file))

    with open(onnx_file, 'rb') as model:
        if not parser.parse(model.read()):
            print('ERROR: Failed to parse the ONNX file.')
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None

    config.set_flag(trt.BuilderFlag.INT8)
    
    setDynamicRange(network, json_file)

    serialized_engine = builder.build_serialized_network(network, config)
    runtime = trt.Runtime(TRT_LOGGER)
    engine = runtime.deserialize_cuda_engine(serialized_engine)

    with open(engine_file, "wb") as f:
        f.write(engine.serialize())
        

if __name__ == '__main__':
    # Add plugins if needed
    # import ctypes
    # ctypes.CDLL("libmmdeploy_tensorrt_ops.so")
    parser = argparse.ArgumentParser(description='Writing qparams to onnx to convert tensorrt engine.')
    parser.add_argument('--onnx', type=str, default=None)
    parser.add_argument('--qparam_json', type=str, default=None)
    parser.add_argument('--engine', type=str, default=None)
    arg = parser.parse_args()
    arg.onnx = ONNX_PATH
    arg.qparam_json = PPQ_JSON_PATH
    arg.engine = ENGINE_OUTPUT_PATH

    build_engine(arg.onnx, arg.qparam_json, arg.engine)
    print("\033[1;32mgenerate %s\033[0m" % arg.engine)


